import os
from typing import List, Optional

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

# 创建FastAPI应用
app = FastAPI(title="DeepSeek-like Chatbot API")

# 允许跨域请求
app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_credentials=True,
	allow_methods=["*"],
	allow_headers=["*"],
)

# 环境变量设置
load_dotenv()
# 调用大语言模型
# 创建模型
model = ChatOpenAI(
	model='qwen-plus',
	api_key=os.getenv("OPENAI_API_KEY"),
	base_url=os.getenv("BASE_URL"),
)

# 获取解析器
parser = StrOutputParser()

# 定义提示模板 - 修改为更符合DeepSeek风格
prompt_template = ChatPromptTemplate.from_messages([
	('system','你是一名AI助手。你乐于助人、友好且知识渊博。请用{language}回答用户的问题。'),
	MessagesPlaceholder(variable_name='chat_history')
])

# 创建处理链
chain = prompt_template | model | parser

# 保存所有用户的聊天记录
store = {}


def get_session_history(session_id: str):
	"""获取或创建指定session_id的聊天历史"""
	if session_id not in store:
		store[session_id] = ChatMessageHistory()
	return store[session_id]


# 创建带消息历史的可运行对象
chat_chain = RunnableWithMessageHistory(
	chain,
	get_session_history,
	input_messages_key="chat_history"
)


# Pydantic模型定义
class ChatMessage(BaseModel):
	content: str
	role: str  # "user" 或 "assistant"


class ChatRequest(BaseModel):
	message: str
	session_id: Optional[str] = "default_session"
	language: Optional[str] = "中文"


class ChatResponse(BaseModel):
	response: str
	session_id: str


class ChatHistoryResponse(BaseModel):
	messages: List[ChatMessage]
	session_id: str


@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
	"""处理用户聊天请求"""
	try:
		# 使用会话ID获取历史
		config = {'configurable': {'session_id': request.session_id}}

		# 调用模型获取响应
		response = chat_chain.invoke(
			{
				'chat_history': [HumanMessage(content=request.message)],
				'language': request.language
			},
			config=config
		)

		return ChatResponse(response=response, session_id=request.session_id)
	except Exception as e:
		raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")


@app.get("/history/{session_id}", response_model=ChatHistoryResponse)
async def get_chat_history(session_id: str):
	"""获取指定会话的聊天历史"""
	try:
		history = get_session_history(session_id)

		# 转换历史记录格式
		messages = []
		for msg in history.messages:
			if msg.type == "human":
				messages.append(ChatMessage(content=msg.content, role="user"))
			elif msg.type == "ai":
				messages.append(ChatMessage(content=msg.content, role="assistant"))

		return ChatHistoryResponse(messages=messages, session_id=session_id)
	except Exception as e:
		raise HTTPException(status_code=500, detail=f"Error retrieving history: {str(e)}")


@app.delete("/history/{session_id}")
async def clear_chat_history(session_id: str):
	"""清除指定会话的聊天历史"""
	try:
		if session_id in store:
			del store[session_id]
		return {"message": f"History for session {session_id} cleared successfully"}
	except Exception as e:
		raise HTTPException(status_code=500, detail=f"Error clearing history: {str(e)}")


@app.get("/")
async def root():
	return {"message": "DeepSeek-like Chatbot API is running!"}


if __name__ == "__main__":
	import uvicorn

	uvicorn.run(app, host="0.0.0.0", port=8000)
